# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# openUBMC is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

"""异步函数异常处理测试

测试async/await语法中的异常处理，包括：
- try/except/finally 在异步函数中的使用
- 异常在 await 边界间的传播
- finally 块在异步函数中的执行
"""

from gtest import EXPECT_EQ, EXPECT_TRUE, test

from mc.future import sleep

cleanup_called = False
finally_executed = False


# ============================================================================
# 基础异常处理测试
# ============================================================================


@test("AsyncException", "BasicTryExcept")
async def test_basic_try_except():
    def normal_func():
        return 42

    async def async_func_that_throws() -> int:
        raise ValueError("test error")
        return 42

    result = 0
    try:
        result = await async_func_that_throws()
    except ValueError:
        result = -1

    EXPECT_EQ(normal_func(), 42)
    EXPECT_EQ(result, -1)


@test("AsyncException", "TryExceptMultipleHandlers")
async def test_try_except_multiple_handlers():
    """测试多个 except 子句"""

    async def async_func_that_throws(exception_type: str) -> int:
        if exception_type == "ValueError":
            raise ValueError("test error")
        elif exception_type == "TypeError":
            raise TypeError("test error")
        elif exception_type == "RuntimeError":
            await sleep(10)
            raise RuntimeError("test error")
        return 42

    async def process(exception_type: str) -> int:
        global cleanup_called
        cleanup_called = False
        try:
            return await async_func_that_throws(exception_type)
        except ValueError:
            return 1
        except TypeError:
            return 2
        except RuntimeError:
            return 3
        finally:
            cleanup_called = True

    EXPECT_EQ(await process("ValueError"), 1)
    EXPECT_TRUE(cleanup_called)

    EXPECT_EQ(await process("TypeError"), 2)
    EXPECT_TRUE(cleanup_called)

    EXPECT_EQ(await process("RuntimeError"), 3)
    EXPECT_TRUE(cleanup_called)

    EXPECT_EQ(await process("UnknownError"), 42)
    EXPECT_TRUE(cleanup_called)


@test("AsyncException", "TryExceptBareExcept")
async def test_try_except_bare_except():
    """测试 bare except（捕获所有异常）"""

    async def async_func_that_throws(v: int) -> int:
        if v == 1:
            raise ValueError("test error")
        elif v == 2:
            raise TypeError("test error")
        elif v == 3:
            raise RuntimeError("test error")
        return 42

    async def process(v: int) -> int:
        try:
            return await async_func_that_throws(v)
        except:
            return 99

    EXPECT_EQ(await process(1), 99)
    EXPECT_EQ(await process(2), 99)
    EXPECT_EQ(await process(3), 99)
    EXPECT_EQ(await process(4), 42)


# ============================================================================
# finally 块测试
# ============================================================================


@test("AsyncException", "TryFinally")
async def test_try_finally():
    """测试 try/finally（无 except）"""

    async def async_func() -> int:
        return 42

    async def process() -> int:
        global cleanup_called
        cleanup_called = False
        try:
            return await async_func()
        finally:
            cleanup_called = True

    EXPECT_EQ(await process(), 42)
    EXPECT_TRUE(cleanup_called)


@test("AsyncException", "TryExceptFinally")
async def test_try_except_finally():
    """测试 try/except/finally"""

    async def async_func_that_throws() -> int:
        await sleep(10)
        raise ValueError("test error")
        return 42

    async def process(v: int) -> int:
        global cleanup_called
        cleanup_called = False
        try:
            if v == 1:
                result = await async_func_that_throws()
                return result
            return await async_func_that_throws() + 1
        except ValueError:
            return v
        finally:
            cleanup_called = True

    EXPECT_EQ(await process(1), 1)
    EXPECT_EQ(await process(2), 2)
    EXPECT_TRUE(cleanup_called)


@test("AsyncException", "TryFinallyWithException")
async def test_try_finally_with_exception():
    """测试 try/finally 中异常传播"""

    async def async_func_that_throws() -> int:
        raise ValueError("test error")
        return 42

    async def process() -> int:
        global cleanup_called
        cleanup_called = False
        try:
            return await async_func_that_throws()
        finally:
            cleanup_called = True

    result = 0
    try:
        result = await process()
        result = 1  # 不应该到达这里
    except Exception:
        result = 2

    EXPECT_EQ(result, 2)
    EXPECT_TRUE(cleanup_called)


# ============================================================================
# 异常在 await 边界间的传播测试
# ============================================================================


@test("AsyncException", "ExceptionInNestedAwait")
async def test_exception_in_nested_await():
    """测试嵌套 await 中的异常传播"""

    async def inner_async() -> int:
        await sleep(10)
        raise ValueError("inner error")
        return 10

    async def outer_async() -> int:
        result = await inner_async()
        return result * 2

    async def process() -> int:
        try:
            result = await outer_async()
            return result
        except ValueError:
            return 0

    EXPECT_EQ(await process(), 0)


# ============================================================================
# 异常与控制流测试
# ============================================================================


@test("AsyncException", "ExceptionInLoop")
async def test_exception_in_loop():
    """测试循环中的异常处理"""

    async def async_func(i: int) -> int:
        if i == 3:
            raise ValueError("error at 3")
        return i * 2

    # 将 process 代码提升到测试函数中，因为只调用一次
    total = 0
    for i in range(5):
        try:
            val = await async_func(i)
            total = total + val
        except ValueError:
            break

    # 0*2 + 1*2 + 2*2 = 6，在 i=3 时异常，break
    EXPECT_EQ(total, 6)


@test("AsyncException", "ExceptionInIf")
async def test_exception_in_if():
    """测试 if 分支中的异常处理，验证各种分支场景"""

    async def async_func(condition: bool) -> int:
        if condition:
            raise ValueError("error")
        return 42

    async def process(condition: bool) -> int:
        try:
            if condition:
                result = await async_func(True)
            else:
                result = await async_func(False)
            return result
        except ValueError:
            return 0

    # 测试各种分支场景
    # 场景1：condition=True，then 分支抛出异常，被 except 捕获
    EXPECT_EQ(await process(True), 0)

    # 场景2：condition=False，else 分支正常执行，无异常
    EXPECT_EQ(await process(False), 42)


# ============================================================================
# else 子句测试
# ============================================================================


@test("AsyncException", "TryExceptElse")
async def test_try_except_else():
    """测试 try/except/else，验证各种分支场景"""

    async def async_func(should_throw: bool) -> int:
        if should_throw:
            raise ValueError("error")
        return 42

    async def process(should_throw: bool) -> int:
        try:
            result = await async_func(should_throw)
        except ValueError:
            return 0
        else:
            return result

        return -1

    # 测试各种分支场景
    # 场景1：should_throw=True，try 块抛出异常，执行 except 分支，不执行 else 分支
    EXPECT_EQ(await process(True), 0)

    # 场景2：should_throw=False，try 块无异常，不执行 except 分支，执行 else 分支
    EXPECT_EQ(await process(False), 42)


# ============================================================================
# 循环中的异常处理 - try 块中的控制流（continue）
# ============================================================================


@test("AsyncException", "TryContinueInLoop")
async def test_try_continue_in_loop():
    """测试循环中 try 块里的 continue，finally 仍执行"""

    async def async_func(i: int) -> int:
        return i * 2

    count = 0
    for i in range(5):
        try:
            if i == 2:
                continue
            val = await async_func(i)
            count = count + val
        except ValueError:
            pass
        finally:
            count = count + 10

    # i=0,1,3,4: val +10 = (0+2+6+8) + 10*4 = 16 + 40 = 56
    # i=2: continue, count +10 = 66
    EXPECT_EQ(count, 66)


# ============================================================================
# 循环中的异常处理 - except 块中的控制流
# ============================================================================


@test("AsyncException", "ExceptBreakInLoop")
async def test_except_break_in_loop():
    """测试循环中 except 块里的 break，finally 仍执行"""

    async def async_func(i: int) -> int:
        if i >= 5:
            raise ValueError("error")
        return i * 2

    count = 0
    for i in range(10):
        try:
            val = await async_func(i)
            count = count + val
        except ValueError:
            count = count + 100
            if i == 5:
                break
        finally:
            count = count + 10

    # i=0,1,2,3,4: val +10 = (0+2+4+6+8) + 10*5 = 20 + 50 = 70
    # i=5: raise, count +100 +10 = 180, then break
    EXPECT_EQ(count, 180)


@test("AsyncException", "ExceptContinueInLoop")
async def test_except_continue_in_loop():
    """测试循环中 except 块里的 continue，finally 仍执行"""

    async def async_func(i: int) -> int:
        if i >= 2:
            raise ValueError("error")
        return i * 2

    count = 0
    for i in range(5):
        try:
            val = await async_func(i)
            count = count + val
        except ValueError:
            if i == 3:
                continue
            count = count + 100
        finally:
            count = count + 10

    # i=0,1: val +10 = (0+2) + 10*2 = 2 + 20 = 22
    # i=2: raise, count +100 +10 = 132
    # i=3: raise, continue, count +10 = 142
    # i=4: raise, count +100 +10 = 252
    EXPECT_EQ(count, 252)


# ============================================================================
# 循环中的异常处理 - else 块中的控制流
# ============================================================================


@test("AsyncException", "ElseBreakInLoop")
async def test_else_break_in_loop():
    """测试循环中 else 块里的 break，finally 仍执行"""

    async def async_func(i: int) -> int:
        return i

    count = 0
    for i in range(10):
        try:
            val = await async_func(i)
        except ValueError:
            pass
        else:
            if i == 5:
                break
        finally:
            count = count + 1

    # break 时 i=5，所以循环执行了 6 次（0-5），count 应该是 6
    EXPECT_EQ(count, 6)


@test("AsyncException", "ElseContinueInLoop")
async def test_else_continue_in_loop():
    """测试循环中 else 块里的 continue，finally 仍执行"""

    async def async_func(i: int) -> int:
        return i

    count = 0
    total = 0
    for i in range(5):
        try:
            val = await async_func(i)
        except ValueError:
            pass
        else:
            if i == 2:
                continue
            total = total + val
        finally:
            count = count + 1

    # count=5, total=0+1+3+4=8
    EXPECT_EQ(count, 5)
    EXPECT_EQ(total, 8)


# ============================================================================
# finally 覆盖控制流 - return
# ============================================================================


@test("AsyncException", "TryReturnWithFinally")
async def test_try_return_with_finally():
    """测试 try 中 return，finally 仍执行"""

    async def async_func() -> int:
        return 10

    async def process() -> int:
        global finally_executed
        finally_executed = False
        try:
            result = await async_func()
            return result
        finally:
            finally_executed = True

    result = await process()
    EXPECT_EQ(result, 10)
    EXPECT_TRUE(finally_executed)


@test("AsyncException", "ExceptReturnWithFinally")
async def test_except_return_with_finally():
    """测试 except 中 return，finally 仍执行"""

    async def async_func() -> int:
        raise ValueError("error")

    async def process() -> int:
        global finally_executed
        finally_executed = False
        try:
            await async_func()
            return 0
        except ValueError:
            return 20
        finally:
            finally_executed = True

    result = await process()
    EXPECT_EQ(result, 20)
    EXPECT_TRUE(finally_executed)


@test("AsyncException", "ElseReturnWithFinally")
async def test_else_return_with_finally():
    """测试 else 块中的 return，finally 仍执行"""

    async def async_func() -> int:
        return 42

    async def process() -> int:
        global finally_executed
        finally_executed = False
        try:
            result = await async_func()
        except ValueError:
            return 0
        else:
            return result
        finally:
            finally_executed = True

    result = await process()
    EXPECT_EQ(result, 42)
    EXPECT_TRUE(finally_executed)


@test("AsyncException", "FinallyReturnOverrideTry")
async def test_finally_return_override_try():
    """测试 finally 中的 return 覆盖 try 中的 return"""

    async def async_func() -> int:
        return 10

    async def process() -> int:
        try:
            result = await async_func()
            return result
        finally:
            return 20  # type: ignore

    result = await process()
    EXPECT_EQ(result, 20)


@test("AsyncException", "FinallyReturnOverrideExcept")
async def test_finally_return_override_except():
    """测试 finally 中的 return 覆盖 except 中的 return"""

    async def async_func() -> int:
        raise ValueError("error")

    async def process() -> int:
        try:
            await async_func()
            return 0
        except ValueError:
            return 30
        finally:
            return 40  # type: ignore

    result = await process()
    EXPECT_EQ(result, 40)


@test("AsyncException", "FinallyReturnOverrideElse")
async def test_finally_return_override_else():
    """测试 finally 中的 return 覆盖 else 中的 return"""

    async def async_func() -> int:
        return 42

    async def process() -> int:
        try:
            result = await async_func()
        except ValueError:
            return 50
        else:
            return 60
        finally:
            return 70  # type: ignore

    result = await process()
    EXPECT_EQ(result, 70)


# ============================================================================
# finally 覆盖控制流 - break（循环中）
# ============================================================================


@test("AsyncException", "FinallyBreakOverrideTry")
async def test_finally_break_override_try():
    """测试 finally 中的 break 覆盖 try 中的 break"""

    async def async_func(i: int) -> int:
        return i

    count = 0
    for i in range(10):
        try:
            val = await async_func(i)
            count = count + val
            if i == 3:
                break
        finally:
            count = count + 1
            if i == 7:
                break  # type: ignore

    # i=0,1,2,3: count +i +1 = (0+1+2+3) + 4 = 10
    # try 中 i=3 时 break，finally 执行后退出，count=10
    EXPECT_EQ(count, 10)


@test("AsyncException", "FinallyBreakDirect")
async def test_finally_break_direct():
    """测试 finally 中直接 break"""

    async def async_func(i: int) -> int:
        return i

    count = 0
    for i in range(10):
        try:
            val = await async_func(i)
            count = count + val
        finally:
            if i == 3:
                break  # type: ignore

    # i=0,1,2,3: count +i = 0+1+2+3 = 6
    EXPECT_EQ(count, 6)


@test("AsyncException", "FinallyBreakOverrideExceptContinue")
async def test_finally_break_override_except_continue():
    """测试 finally 中的 break 覆盖 except 中的 continue"""

    async def async_func(i: int) -> int:
        if i >= 5:
            raise ValueError("error")
        return i

    count = 0
    total = 0
    for i in range(10):
        try:
            val = await async_func(i)
            count = count + val
        except ValueError:
            total = total + 1
            continue
        finally:
            if i == 7:
                break  # type: ignore

    # i=0,1,2,3,4: count +i = 0+1+2+3+4 = 10
    # i=5,6: raise, total +1 = 2, continue
    # i=7: raise, total +1 = 3, continue 被 finally break 覆盖
    EXPECT_EQ(count, 10)
    EXPECT_EQ(total, 3)


# ============================================================================
# finally 覆盖控制流 - continue（循环中）
# ============================================================================


@test("AsyncException", "FinallyContinueDirect")
async def test_finally_continue_direct():
    """测试 finally 中直接 continue"""

    async def async_func(i: int) -> int:
        return i

    count = 0
    total = 0
    for i in range(5):
        try:
            val = await async_func(i)
            count = count + val
        finally:
            if i == 2:
                continue  # type: ignore
            total = total + i

    # i=0,1,2,3,4: count +i = 0+1+2+3+4 = 10
    # i=0,1,3,4: total +i = 0+1+3+4 = 8（跳过 i=2）
    EXPECT_EQ(count, 10)
    EXPECT_EQ(total, 8)


@test("AsyncException", "FinallyContinueOverrideTryBreak")
async def test_finally_continue_override_try_break():
    """测试 finally 中的 continue 覆盖 try 中的 break"""

    async def async_func(i: int) -> int:
        return i

    count = 0
    for i in range(10):
        try:
            i += 1
            i -= 1
            val = await async_func(i)
            count = count + val
            if i == 3:
                break
        finally:
            if i == 3:
                continue  # type: ignore

    # i=0,1,2: count +i = 0+1+2 = 3
    # i=3: count +i = 6, break 被 finally continue 覆盖，继续循环
    # i=4,5,6,7,8,9: count +i = 6+4+5+6+7+8+9 = 45
    EXPECT_EQ(count, 45)


@test("AsyncException", "FinallyContinueOverrideExceptBreak")
async def test_finally_continue_override_except_break():
    """测试 finally 中的 continue 覆盖 except 中的 break"""

    async def async_func(i: int) -> int:
        if i >= 5:
            raise ValueError("error")
        return i

    count = 0
    for i in range(10):
        try:
            val = await async_func(i)
            count = count + val
        except ValueError:
            count = count + 100
            if i == 5:
                break
        finally:
            if i == 5:
                continue  # type: ignore

    # i=0,1,2,3,4: count +i = 0+1+2+3+4 = 10
    # i=5: raise, count +100 = 110, break 被 finally continue 覆盖
    # i=6,7,8,9: raise, count +100 = 510
    EXPECT_EQ(count, 510)


@test("AsyncException", "FinallyReturnOverrideTryBreak")
async def test_finally_return_override_try_break():
    """测试 finally 中的 return 覆盖 try 中的 break"""

    async def async_func(i: int) -> int:
        return i

    async def process() -> int:
        count = 0
        for i in range(10):
            try:
                val = await async_func(i)
                count = count + val
                if i == 3:
                    break
            finally:
                if i == 3:
                    return count  # type: ignore
        return count

    result = await process()
    # i=0,1,2,3: count +i = 0+1+2+3 = 6
    EXPECT_EQ(result, 6)


# ============================================================================
# 嵌套异常测试
# ============================================================================


@test("AsyncException", "NestedTryExcept")
async def test_nested_try_except():
    """测试嵌套的 try/except"""

    async def inner_async() -> int:
        raise ValueError("inner error")

    outer_caught = False
    inner_caught = False

    try:
        try:
            await inner_async()
        except ValueError:
            inner_caught = True
            raise RuntimeError("outer error")
    except RuntimeError:
        outer_caught = True

    EXPECT_TRUE(inner_caught)
    EXPECT_TRUE(outer_caught)


finally_count = 0


@test("AsyncException", "NestedFinallyReturn")
async def test_nested_finally_return():
    """测试嵌套 finally 的执行顺序"""

    async def async_func() -> int:
        return 10

    async def process() -> int:
        global finally_count
        try:
            try:
                result = await async_func()
                return result
            finally:
                finally_count = finally_count + 1
        finally:
            finally_count = finally_count + 10

    result = await process()
    EXPECT_EQ(result, 10)
    EXPECT_EQ(finally_count, 11)


@test("AsyncException", "NestedFinallyException")
async def test_nested_finally_exception():
    """测试嵌套 finally 的异常处理"""

    async def async_func() -> int:
        raise ValueError("inner")

    result = 0
    try:
        try:
            await async_func()
        except ValueError:
            result = result + 1
        finally:
            result = result + 10
    except RuntimeError:
        result = result + 100
    finally:
        result = result + 1000

    # inner exception 被捕获 (result=1) → inner finally (result=11) → outer finally (result=1011)
    EXPECT_EQ(result, 1011)


@test("AsyncException", "NestedFinallyWithCodeAfter")
async def test_nested_finally_with_code_after():
    """测试内层 finally 之后的正常代码执行时机"""

    async def async_func() -> int:
        raise ValueError("inner")

    result = 0
    try:
        try:
            await async_func()
        except ValueError:
            result = result + 1
        finally:
            result = result + 10

        # 内层 try/except/finally 之后的正常代码
        result = result + 100
    finally:
        result = result + 1000

    # inner except: result = 1
    # inner finally: result = 11
    # 正常代码: result = 111
    # outer finally: result = 1111
    EXPECT_EQ(result, 1111)


# ============================================================================
# 异常重新抛出和包装测试
# ============================================================================


@test("AsyncException", "BareRaise")
async def test_bare_raise():
    """测试 bare raise（重新抛出当前异常）"""

    async def async_func() -> int:
        raise ValueError("test")

    exception_caught = False

    try:
        try:
            await async_func()
        except ValueError:
            raise  # bare raise，重新抛出当前异常
    except ValueError:
        exception_caught = True

    EXPECT_TRUE(exception_caught)


@test("AsyncException", "ExceptRaiseWithFinally")
async def test_except_raise_with_finally():
    """测试 except 块中 raise 新异常时 finally 的执行"""

    async def async_func() -> int:
        raise ValueError("inner")

    global finally_executed
    finally_executed = False
    exception_caught = False

    try:
        try:
            await async_func()
        except ValueError:
            finally_executed = True
            raise RuntimeError("outer")
        finally:
            finally_executed = True
    except RuntimeError:
        exception_caught = True

    EXPECT_TRUE(exception_caught)
    EXPECT_TRUE(finally_executed)


@test("AsyncException", "NestedTryWithStatementsEverywhere")
async def test_nested_try_with_statements_everywhere():
    """测试嵌套 try 块在各个位置的普通语句处理"""

    async def async_func(should_throw: bool) -> int:
        if should_throw:
            raise ValueError("inner error")
        return 100

    result = 0
    trace = ""

    # 外层 try 前的语句
    result = result + 1  # result = 1
    trace = trace + "A"

    try:
        # 外层 try 开始后、内层 try 前的语句
        result = result + 2  # result = 3
        trace = trace + "B"

        try:
            # 内层 try 开始后、await 前的语句
            result = result + 4  # result = 7
            trace = trace + "C"

            val = await async_func(True)

            # await 后的语句（异常时不执行）
            result = result + 8
            trace = trace + "D"

        except ValueError:
            # 内层 except 中的语句
            result = result + 16  # result = 23
            trace = trace + "E"
        finally:
            # 内层 finally 中的语句
            result = result + 32  # result = 55
            trace = trace + "F"

        # 内层 try-except-finally 后的语句
        result = result + 64  # result = 119
        trace = trace + "G"

    finally:
        # 外层 finally 中的语句
        result = result + 128  # result = 247
        trace = trace + "H"

    # 外层 try-finally 后的语句
    result = result + 256  # result = 503
    trace = trace + "I"

    # 验证结果：1 + 2 + 4 + 16 + 32 + 64 + 128 + 256 = 503
    EXPECT_EQ(result, 503)
    # 验证执行顺序：A(外层try前) B(外层try内) C(内层try内) E(内层except) F(内层finally) G(内层后) H(外层finally) I(外层后)
    EXPECT_EQ(trace, "ABCEFGHI")


@test("AsyncException", "ElseRaiseWithFinally")
async def test_else_raise_with_finally():
    """测试 else 块中 raise 时 finally 的执行"""

    async def async_func() -> int:
        return 42

    global finally_executed
    finally_executed = False
    exception_caught = False

    try:
        try:
            result = await async_func()
        except ValueError:
            pass
        else:
            raise RuntimeError("from else")
        finally:
            finally_executed = True
    except RuntimeError:
        exception_caught = True

    EXPECT_TRUE(exception_caught)
    EXPECT_TRUE(finally_executed)
